#pragma once

#include "updater_common.h"

namespace mdtx::online::updater
{

    // arithmetic s2
    template <bool kahan>
    struct s2_updater
    {
        size_t n;
        double m, s2;
        accumulator<kahan> acc_m, acc_s2;

        s2_updater()
            : n(0),
              m(0.0),
              s2(0.0)
        {
        }

        void insert(double x)
        {
            double d = x - m;
            n += 1;
            acc_m(m, d / n);
            acc_s2(s2, (x - m) * d);
        }

        void remove(double x)
        {
            double d = m - x;
            n -= 1;
            acc_m(m, d / n);
            acc_s2(s2, (x - m) * d);
        }
    };

    // weighted s2
    // West, D. H. D. (1979). Updating mean and variance estimates: An improved method. Communications of the ACM, 22(9), 532-535.
    template <bool kahan>
    struct s2w_updater
    {
        double wsum, m, s2;
        accumulator<kahan> acc_w, acc_m, acc_s2;

        s2w_updater()
            : wsum(0.0),
              m(0.0),
              s2(0.0)
        {
        }

        void insert(double x, double w)
        {
            double d = x - m;
            acc_w(wsum, w);
            acc_m(m, (w / wsum) * d);
            acc_s2(s2, w * (x - m) * d);
        }

        void remove(double x, double w)
        {
            double d = m - x;
            acc_w(wsum, -w);
            acc_m(m, (w / wsum) * d);
            acc_s2(s2, w * (x - m) * d);
        }
    };

    // Sum of w^2 is also calculated for reliability weights
    template <bool kahan>
    struct s2wr_updater
    {
        double wsum, w2sum, m, s2;
        accumulator<kahan> acc_w, acc_w2, acc_m, acc_s2;

        s2wr_updater()
            : wsum(0.0),
              w2sum(0.0),
              m(0.0),
              s2(0.0)
        {
        }

        void insert(double x, double w)
        {
            double d = x - m;
            acc_w(wsum, w);
            acc_w2(w2sum, w * w);
            acc_m(m, (w / wsum) * d);
            acc_s2(s2, w * (x - m) * d);
        }

        void remove(double x, double w)
        {
            double d = m - x;
            acc_w(wsum, -w);
            acc_w2(w2sum, -w * w);
            acc_m(m, (w / wsum) * d);
            acc_s2(s2, w * (x - m) * d);
        }
    };

    // exponentially weighted s2
    template <bool kahan>
    struct s2ew_updater
    {
        bool init;
        double a, b, m, s2;
        accumulator<kahan> acc_m;

        explicit s2ew_updater(double alpha)
            : init(false),
              a(alpha),
              b(1.0 - a),
              m(0.0),
              s2(0.0)
        {
        }

        void insert(double x)
        {
            if (init)
            {
                double d = x - m;
                double i = a * d;
                acc_m(m, i);
                s2 = b * (s2 + i * d);
            }
            else
            {
                m = x;
                init = true;
            }
        }

        void remove(double x)
        {
            // we don't remove from EMA
            x = x;
        }
    };
}
